import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import rasterio


def get_background(tif_file):
    with rasterio.open(tif_file) as src:
        img = src.read([1, 2, 3])  # 取 RGB 三个波段
        print(img.shape)
        print(img[0].min(),img[0].max())
        print(img[1].min(),img[1].max())
        print(img[2].min(),img[2].max())
        bounds = src.bounds        # 获取地理范围
        extent = [bounds.left, bounds.right, bounds.bottom, bounds.top]
        img_stretched = contrast_stretch(img, 4, 96)
        img_stretched=img_stretched ** 0.6
    return img_stretched,extent

# from rasterio.plot import show
def contrast_stretch(img, low=2, high=98):
    """对多波段影像做百分位拉伸"""
    out = np.zeros_like(img, dtype=np.float32)
    for b in range(img.shape[0]):  # 每个波段单独拉伸
        p_low, p_high = np.percentile(img[b], (low, high))
        out[b] = np.clip((img[b] - p_low) / (p_high - p_low), 0, 1)
    return out

def get_trackfile(trackfile):
    ff = np.load(trackfile)
    clon = ff[:,1].astype(float)
    clat = ff[:,2].astype(float)
    return clon,clat

def get_trackfiles(trackfiles):
    nn=len(trackfiles)
    clons=[]
    clats=[]
    for i in range(nn):
        clon,clat = get_trackfile(trackfiles[i])
        clons.append(clon)
        clats.append(clat)
    return clons,clats


def final_plot(tif_file,trackfiles):
    img_stretched,extent = get_background(tif_file)
    clons,clats = get_trackfiles(trackfiles)
    nn=len(clons)
    # === 作图 ===
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent(extent, ccrs.PlateCarree())
    ax.imshow(np.transpose(img_stretched, (1, 2, 0)), extent=extent, origin="upper")
    for i in range(nn):
        ax.plot(clons[i], clats[i], 'r-', transform=ccrs.PlateCarree(), linewidth=0.5, label="")

    # 经纬度网格
    gl = ax.gridlines(draw_labels=True, linestyle="--", alpha=0.5)
    gl.top_labels = False
    gl.right_labels = False

    #ax.legend()
    plt.show()
    plt.savefig('202_center_track.png',dpi=1000)

# === 轨迹数据 ===
track_files =[r'202_center_track.npy']
# === 读取 GeoTIFF ===
tif_file = r"E:\akesu_radar\cheng\Xinjiang_RGB.tif"   # 你导出的 GeoTIFF 文件
final_plot(tif_file,track_files)